from yolov5_face_detection import yolov5_face_detector
from yolov3_plate_detection import yolov3_license_plate_detector
import cv2
import numpy as np
import os
import json
from time import time
import argparse


def rel2abspath(path):
    # 相对路径转化为绝对路径
    # 输出路径写入文件相对路径报错
    if os.path.isabs(path):
        pass
    else:
        path = os.path.join(os.getcwd(), path)
    return path

def cv_img_read(file_path):
    # cv2.imread读取中文路径出错，替代
    cv_img = cv2.imdecode(np.fromfile(file_path, dtype=np.uint8), cv2.IMREAD_COLOR)
    return cv_img

def cv_img_write(output_path, image):
    # cv2.imwrite写入中文路径出错，替代
    cv2.imencode('.jpg', image)[1].tofile(output_path)

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--input-path", type=str, required=True)
    parser.add_argument("--output-path", type=str, required=True)
    parser.add_argument("--work-mode", type=str, choices=["face_detect", "face_mosaic", "data_mask", 
                        "detect_both", "plate_mosaic", "plate_detect"], default="data_mask")
    parser.add_argument("--mosaic-level", type=int, default=3)
    parser.add_argument("--face-yolo-mode", choices=['yolov5s', 'yolov5m', 'yolov5l'], type=str,
                        default="yolov5s")
    parser.add_argument("--face-yolo-conf-thresh", type=float, default=0.5)   
    parser.add_argument("--face-yolo-obj-thresh", type=float, default=0.5)
    parser.add_argument("--face-yolo-nms-iou", type=float, default=0.5)
    parser.add_argument("--plate-yolo-conf-thresh", type=float, default=0.5)
    parser.add_argument("--plate-yolo-nms-iou", type=float, default=0.5)

    args = parser.parse_args()

    return args

class Data_Mask():
    def __init__(self, args):
        self.work_mode = args.work_mode

        self.input_path = args.input_path
        self.output_path = rel2abspath(args.output_path)
        self._suffix_name_init()
        
        self.mosaic_level = args.mosaic_level

        self.face_yolo_mode = args.face_yolo_mode
        self.face_yolo_conf_thresh = args.face_yolo_conf_thresh
        self.face_yolo_obj_thresh = args.face_yolo_obj_thresh
        self.face_yolo_nms_iou = args.face_yolo_nms_iou
        self.face_detector = yolov5_face_detector(yolo_type=self.face_yolo_mode,
                                                  confThreshold=self.face_yolo_conf_thresh,
                                                  nmsThreshold=self.face_yolo_nms_iou,
                                                  objThreshold=self.face_yolo_obj_thresh)

        self.plate_yolo_conf_thresh = args.plate_yolo_conf_thresh
        self.plate_yolo_nms_iou = args.plate_yolo_nms_iou
        self.plate_detector = yolov3_license_plate_detector(confThreshold=self.plate_yolo_conf_thresh,
                                                            nmsThreshold=self.plate_yolo_nms_iou)
        
    def _suffix_name_init(self):
        work_modes = ["face_detect", "face_mosaic", "data_mask", "detect_both",
                      "plate_mosaic", "plate_detect"]
        suffix_name_idx = [self.work_mode == opt for opt in work_modes]
        self.suffix_name = '_' + work_modes[suffix_name_idx.index(True)]

    def work_process(self, srcimg):
        face_num = 0
        plate_num = 0
        if self.work_mode in ["face_detect", "face_mosaic"]:
            face_dets = self.face_detector.detect(srcimg)
            if self.work_mode == "face_detect":
                result_file, face_num, face_boxes = self.face_detector.postprocess(srcimg, face_dets, mark='origin')
            elif self.work_mode == "face_mosaic":
                mosaic_level_in = (self.mosaic_level + 4) * 0.1
                result_file, face_num, face_boxes = self.face_detector.postprocess(srcimg, face_dets, mark='mosaic',
                                                                       mosaic_level=mosaic_level_in)
        elif self.work_mode in ["plate_mosaic", "plate_detect"]:
            if self.work_mode == "plate_detect":
                result_file, plate_num, plate_boxes = self.plate_detector.rectangle(srcimg, mode='origin')
            elif self.work_mode == "plate_mosaic":
                mosaic_level_in = (self.mosaic_level + 4) * 0.1
                result_file, plate_num, plate_boxes = self.plate_detector.rectangle(srcimg, mode='mosaic',
                                                                       mosaic_level=mosaic_level_in)
        elif self.work_mode == "data_mask":
            face_dets = self.face_detector.detect(srcimg)
            mosaic_level_in = (self.mosaic_level + 4) * 0.1
            result_file, face_num, face_boxes = self.face_detector.postprocess(srcimg, face_dets, mark='mosaic',
                                                                   mosaic_level=mosaic_level_in)
            result_file, plate_num, plate_boxes = self.plate_detector.rectangle(result_file, mode='mosaic',
                                                                   mosaic_level=mosaic_level_in)
        elif self.work_mode == "detect_both":
            face_dets = self.face_detector.detect(srcimg)
            result_file, face_num, face_boxes = self.face_detector.postprocess(srcimg, face_dets, mark='origin')
            result_file, plate_num, plate_boxes = self.plate_detector.rectangle(result_file, mode='origin')

        self.result_file = result_file
        return result_file, face_num, plate_num, face_boxes, plate_boxes

    def forward_single(self):
        srcimg = cv_img_read(self.input_path)
        if srcimg is not None:
            result_file, face_num, plate_num, _ = self.work_process(srcimg=srcimg)
            image_name = os.path.basename(self.input_path)
            image_new_name = image_name.split('.')[0] + self.suffix_name + '.jpg'
            result_image_path = os.path.join(self.output_path, image_new_name)
            cv_img_write(result_image_path, result_file)

            print(f"\n{face_num} faces detected in the image")
            print(f"\n{plate_num} license plates detected in the image")
        else:
            print(f"\nPlease input image file!")

    def forward_batch(self):
        if self.input_path is not None and self.output_path is not None:
            # traverse every image in input path
            print("\nBatch Process Start...")
            images_list = os.listdir(self.input_path)
            files_len = len(images_list)
            start_time = time()
            not_image_num = 0
            dir_num = 0
            process_face_image_num = 0
            process_plate_image_num = 0
            process_image_num = 0
            out_dict = {}
            for idx, img_name in enumerate(images_list):
                # get image
                img_path = os.path.join(self.input_path, img_name)
                if os.path.isfile(img_path):
                    srcimg = cv_img_read(img_path)
                    if srcimg is not None:
                        img_new_name = img_name.split('.')[0] + self.suffix_name + '.jpg'
                        srcimg, face_num, plate_num, face_boxes, plate_boxes = self.work_process(srcimg=srcimg)
                        if face_num >= 1 or plate_num >= 1:
                            if face_num >= 1:
                                process_face_image_num += 1
                            if plate_num >= 1:
                                process_plate_image_num += 1
                            process_image_num += 1
                            cv_img_write(os.path.join(self.output_path, img_new_name), srcimg)
                            out_dict[f'{img_new_name}'] = f'face:{face_boxes}  plate:{plate_boxes}'
                    else:
                        not_image_num += 1
                else:
                    dir_num += 1
                print('\r'+'▋'*(idx+1)+'当前进度：{:.2f}%'.format((idx+1)/files_len*100),end='')
            json_write = json.dumps(out_dict, indent=1)
            with open(f'{self.output_path}/result.json', 'w') as jw:
                jw.write(json_write)
            end_time = time()
            duration = end_time - start_time
            print("\nBatch Process Success.")
            print(f"\nImage(with object) num: {process_image_num}")
            print(f"Image(with face) num: {process_face_image_num}")
            print(f"Image(with license plate) num: {process_plate_image_num}")
            print(f"\nImage total num: {files_len - not_image_num - dir_num}")
            print(f"Directory num: {dir_num}")
            print(f"Other file num: {not_image_num}")
            print(f"Total num: {files_len}")
            print(f"\nTime spend: {duration:.3f}s")
        else:
            print("\nWARNING: Please select set both input path and output path!")


if __name__ == "__main__":
    args = parse_args()
    worker = Data_Mask(args)
    if os.path.isdir(args.input_path):
        worker.forward_batch()
    else:
        worker.forward_single()
